﻿using System;
using System.Collections.Generic;
using System.Collections.Specialized;
using System.Configuration;
using System.Configuration.Provider;
using System.Data;
using System.Diagnostics;
using System.IO;
using System.Web;
using System.Web.Configuration;
using System.Web.SessionState;

using Aerospike.Client;

/*
This session state store provider supasPorts the following schema:

  CREATE TABLE Sessions
  (
    SessionId       Text(80)  NOT NULL,
    ApplicationName Text(255) NOT NULL,
    Created         DateTime  NOT NULL,
    Expires         DateTime  NOT NULL,
    LockDate        DateTime  NOT NULL,
    LockId          Integer   NOT NULL,
    Timeout         Integer   NOT NULL,
    Locked          YesNo     NOT NULL,
    SessionItems    Memo,
    Flags           Integer   NOT NULL,

    CONSTRAINT PKSessions PRIMARY KEY (SessionId, ApplicationName)
  )

This session state store provider does not automatically clean up 
expired session item data. It is recommended
that you periodically delete expired session information from the
data store with the following code (where 'conn' is the OdbcConnection
for the session state store provider):

  string commandString = "DELETE FROM Sessions WHERE Expires < ?";
  OdbcConnection conn = new OdbcConnection(connectionString);
  OdbcCommand cmd = new OdbcCommand(commandString, conn);
  cmd.Parameters.Add("@Expires", OdbcType.DateTime).Value = DateTime.UtcNow;
  conn.Open();
  cmd.ExecuteNonQuery();
  conn.Close();

*/

namespace RepX.Web.SessionState
{

    public sealed class AerospikeSessionStateStore : SessionStateStoreProviderBase
    {
        private SessionStateSection pConfig = null;

        private const int MAX_DATA_SIZE = 128 * 1024 - 128;

        private Host[] asHosts;
        private int asPort = 3000;
        private string asNamespace = "AerospikeSessionState";
        private string sessionSet = "Session";
        private string sessionDataSet = "SessionData";
        private bool asCompress = true;

        private string eventSource = "AerospikeSessionStateStore";
        private string eventLog = "Application";
        private string exceptionMessage = "An exception occurred. Please contact your administrator.";
        private string pApplicationName;

        private ClientPolicy clientPolicy = new ClientPolicy();
        private WritePolicy writePolicy = new WritePolicy();


        public AerospikeSessionStateStore()
        {
            // Default Policies
            writePolicy.recordExistsAction = RecordExistsAction.UPDATE;
        }

        //
        // If false, exceptions are thrown to the caller. If true,
        // exceptions are written to the event log.
        //

        private bool pWriteExceptionsToEventLog = false;

        public bool WriteExceptionsToEventLog
        {
            get { return pWriteExceptionsToEventLog; }
            set { pWriteExceptionsToEventLog = value; }
        }


        //
        // The ApplicationName property is used to differentiate sessions
        // in the data source by application.
        //

        public string ApplicationName
        {
            get { return pApplicationName; }
        }


        public override void Initialize(string name, NameValueCollection config)
        {
            //
            // Initialize values from web.config.
            //

            if (config == null) throw new ArgumentNullException("config");

            if (name == null || name.Length == 0) name = "AerospikeSessionStateStore";

            if (String.IsNullOrEmpty(config["description"]))
            {
                config.Remove("description");
                config.Add("description", "Sample ODBC Session State Store provider");
            }

            // Initialize the abstract base class.
            base.Initialize(name, config);


            //
            // Initialize the ApplicationName property.
            //

            pApplicationName = System.Web.Hosting.HostingEnvironment.ApplicationVirtualPath;


            //
            // Get <sessionState> configuration element.
            //

            System.Configuration.Configuration cfg = WebConfigurationManager.OpenWebConfiguration(ApplicationName);
            pConfig = (SessionStateSection)cfg.GetSection("system.web/sessionState");



            //
            // Load Connections
            //

            if (config["Compress"] != null) asCompress = Convert.ToBoolean(config["Compress"]);
            if (config["Namespace"] != null) asNamespace = config["Namespace"];
            if (config["Port"] != null) asPort = Convert.ToInt32(config["Port"]);
            if (config["Host"] != null)
            {
                List<Host> h = new List<Host>();

                string[] hosts = config["Host"].Split(new char[] { ',' });

                // Create Multiple Hosts
                foreach (string host in hosts)
                {
                    h.Add(new Host(host, asPort));
                }

                asHosts = h.ToArray();
            }


            //
            // Initialize WriteExceptionsToEventLog
            //

            pWriteExceptionsToEventLog = false;

            if (config["writeExceptionsToEventLog"] != null)
            {
                if (config["writeExceptionsToEventLog"].ToUpper() == "TRUE") pWriteExceptionsToEventLog = true;
            }
        }


        //
        // SessionStateStoreProviderBase members
        //

        public override void Dispose()
        {
        }


        //
        // SessionStateProviderBase.SetItemExpireCallback
        //

        public override bool SetItemExpireCallback(SessionStateItemExpireCallback expireCallback)
        {
            return false;
        }


        //
        // SessionStateProviderBase.SetAndReleaseItemExclusive
        //

        public override void SetAndReleaseItemExclusive(HttpContext context, string id, SessionStateStoreData item, object lockId, bool newItem)
        {
            // Serialize the SessionStateItemCollection as a string.
            byte[] sessItems = Serialize((SessionStateItemCollection)item.Items);

            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);

            try
            {
                if (newItem)
                {
                    // Delete Expired Session
                    Record record = client.Get(new Policy(), key, new string[] { "AN", "Expires" });
                    if (record != null &&
                        Convert.ToString(record.bins["AN"]) == ApplicationName &&
                        Convert.ToDateTimeNullable(record.bins["Expires"]) <= DateTime.UtcNow)
                    {
                        client.Delete(writePolicy, key);
                    }

                    List<Bin> bins = new List<Bin>()
                    {
                        new Bin("SessionId", id),
                        new Bin("AN", ApplicationName),
                        new Bin("Created", DateTime.UtcNow),
                        new Bin("Expires", DateTime.UtcNow.AddMinutes((Double)item.Timeout)),
                        new Bin("LockDate", DateTime.UtcNow),
                        new Bin("LockId", (Int32)0),
                        new Bin("Timeout", item.Timeout),
                        new Bin("Locked", false),
                        new Bin("Flags", 0),
                        new Bin("SI", (Int32)0),
                    };

                    client.Put(writePolicy, key, bins.ToArray());

                    //bins.AddRange(this.AddData(sessItems, "SI"));

                    // Write Session Data
                    List<byte[]> dataBlock = this.SplitData(sessItems, "");

                    for (int i = 0; i * MAX_DATA_SIZE < sessItems.Length; i++)
                    {
                        Key dataKey = new Key(asNamespace, sessionDataSet, id + ":" + i.ToString());

                        client.Put(writePolicy, dataKey, new Bin("V", dataBlock[i]));
                    }
                }
                else
                {
                    Record record = client.Get(new Policy(), key, new string[] { "AN", "LockId" });
                    if (record != null &&
                        Convert.ToString(record.bins["AN"]) == ApplicationName &&
                        Convert.ToInt32(record.bins["LockId"]) == Convert.ToInt32(lockId))
                    {
                        List<Bin> bins = new List<Bin>()
                        {
                            new Bin("Expires", DateTime.UtcNow.AddMinutes((Double)item.Timeout)),
                            new Bin("Locked", false),
                            new Bin("SI", sessItems.Length),
                        };

                        // Write Session
                        client.Put(writePolicy, key, bins.ToArray());

                        // Write Session Data
                        List<byte[]> dataBlock = this.SplitData(sessItems, "");

                        for (int i = 0; i * MAX_DATA_SIZE < sessItems.Length; i++)
                        {
                            Key dataKey = new Key(asNamespace, sessionDataSet, id + ":" + i.ToString());

                            client.Put(writePolicy, dataKey, new Bin("V", dataBlock[i]));
                        }
                    }
                }
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "SetAndReleaseItemExclusive");
                    throw new ProviderException(exceptionMessage);
                }
                else throw new AerospikeException(string.Format("Session block size exceeded: {0}", sessItems.Length), e);
            }
            finally
            {
                client.Close();
            }
        }


        //
        // SessionStateProviderBase.GetItem
        //

        public override SessionStateStoreData GetItem(HttpContext context, string id, out bool locked, out TimeSpan lockAge, out object lockId, out SessionStateActions actionFlags)
        {
            return GetSessionStoreItem(false, context, id, out locked, out lockAge, out lockId, out actionFlags);
        }


        //
        // SessionStateProviderBase.GetItemExclusive
        //

        public override SessionStateStoreData GetItemExclusive(HttpContext context, string id, out bool locked, out TimeSpan lockAge, out object lockId, out SessionStateActions actionFlags)
        {
            return GetSessionStoreItem(true, context, id, out locked, out lockAge, out lockId, out actionFlags);
        }


        //
        // GetSessionStoreItem is called by both the GetItem and 
        // GetItemExclusive methods. GetSessionStoreItem retrieves the 
        // session data from the data source. If the lockRecord parameter
        // is true (in the case of GetItemExclusive), then GetSessionStoreItem
        // locks the record and sets a new LockId and LockDate.
        //

        private SessionStateStoreData GetSessionStoreItem(bool lockRecord, HttpContext context, string id, out bool locked, out TimeSpan lockAge, out object lockId, out SessionStateActions actionFlags)
        {
            // Initial values for return value and out parameters.
            SessionStateStoreData item = null;
            lockAge = TimeSpan.Zero;
            lockId = null;
            locked = false;
            actionFlags = 0;

            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);


            // DateTime to check if current session item is expired.
            DateTime expires;
            // String to hold serialized SessionStateItemCollection.
            byte[] serializedItems = null;
            // True if a record is found in the database.
            bool foundRecord = false;
            // True if the returned session item is expired and needs to be deleted.
            bool deleteData = false;
            // Timeout value from the data store.
            int timeout = 0;

            try
            {
                // lockRecord is true when called from GetItemExclusive and
                // false when called from GetItem.
                // Obtain a lock if possible. Ignore the record if it is expired.
                if (lockRecord)
                {
                    Record record = client.Get(new Policy(), key, new string[] { "AN", "Locked", "Expires" });

                    if (record != null &&
                        Convert.ToString(record.bins["AN"]) == ApplicationName &&
                        Convert.ToBoolean(record.bins["Locked"]) == false &&
                        Convert.ToDateTime(record.bins["Expires"]) > DateTime.UtcNow)
                    {
                        List<Bin> bins = new List<Bin>()
                        {
                            new Bin("Locked", true),
                            new Bin("LockDate", DateTime.UtcNow),
                        };

                        client.Put(writePolicy, key, bins.ToArray());

                        locked = false;     // The record was updated.
                    }
                    else locked = true;     // No record was updated because the record was locked or not found.
                }

                Record session = client.Get(new Policy(), key); //, new string[] { "AN", "Expires", "SI", "LockId", "LockDate", "Flags", "Timeout", "Locked" });
                if (session != null)
                {
                    expires = Convert.ToDateTime(session.bins["Expires"]);

                    if (expires < DateTime.UtcNow)
                    {
                        // The record was expired. Mark it as not locked.
                        locked = false;

                        // The session was expired. Mark the data for deletion.
                        deleteData = true;
                    }
                    else foundRecord = true;

                    lockId = Convert.ToInt32(session.bins["LockId"]);
                    lockAge = DateTime.UtcNow.Subtract(Convert.ToDateTime(session.bins["LockDate"]));
                    actionFlags = (SessionStateActions)Convert.ToInt32(session.bins["Flags"]);
                    timeout = Convert.ToInt32(session.bins["Timeout"]);

                    //serializedItems = Convert.ToByteArray(session.bins["SI"]);
                    //serializedItems = this.ExtractData(session.bins, "SI");
                    int Size = Convert.ToInt32(session.bins["SI"]);

                    if (!lockRecord) locked = Convert.ToBoolean(session.bins["Locked"]);

                    // Read SessionData
                    List<Key> dataKeys = new List<Key>();
                    for (int i = 0; i * MAX_DATA_SIZE < Size; i++)
                    {
                        dataKeys.Add(new Key(asNamespace, sessionDataSet, id + ":" + i.ToString()));
                    }

                    Record[] sessionData = client.Get(new BatchPolicy(), dataKeys.ToArray());
                    List<byte[]> dataBlock = new List<byte[]>();
                    foreach (Record dataRecord in sessionData)
                    {
                        dataBlock.Add(Convert.ToByteArray(dataRecord.bins["V"]));
                    }

                    serializedItems = this.CombineData(dataBlock, Size);
                }

                // If the returned session item is expired, 
                // delete the record from the data source.
                if (deleteData)
                {
                    client.Delete(writePolicy, key);
                }

                // The record was not found. Ensure that locked is false.
                if (!foundRecord) locked = false;

                // If the record was found and you obtained a lock, then set 
                // the lockId, clear the actionFlags,
                // and create the SessionStateStoreItem to return.
                if (foundRecord && !locked)
                {
                    lockId = (int)lockId + 1;

                    List<Bin> bins = new List<Bin>()
                    {
                        new Bin("LockId", Convert.ToInt32(lockId)),
                        new Bin("Flags", 0),
                    };

                    client.Put(writePolicy, key, bins.ToArray());

                    // If the actionFlags parameter is not InitializeItem, 
                    // deserialize the stored SessionStateItemCollection.
                    if (actionFlags == SessionStateActions.InitializeItem) item = CreateNewStoreData(context, (int)pConfig.Timeout.TotalMinutes);
                    else item = Deserialize(context, serializedItems, timeout);
                }
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "GetSessionStoreItem");
                    throw new ProviderException(exceptionMessage);
                }
                else throw e;
            }
            finally
            {
                client.Close();
            }

            return item;
        }




        //
        // Serialize is called by the SetAndReleaseItemExclusive method to 
        // convert the SessionStateItemCollection into a Base64 string to    
        // be stored in an Access Memo field.
        //

        private byte[] Serialize(SessionStateItemCollection items)
        {
            using (MemoryStream ms = new MemoryStream())
            {
                BinaryWriter writer = new BinaryWriter(ms);

                if (items != null) items.Serialize(writer);

                writer.Close();

                return (this.asCompress ? Convert.Compress(ms.ToArray()) : ms.ToArray());
            }
        }

        //
        // DeSerialize is called by the GetSessionStoreItem method to 
        // convert the Base64 string stored in the Access Memo field to a 
        // SessionStateItemCollection.
        //

        private SessionStateStoreData Deserialize(HttpContext context, byte[] serializedItems, int timeout)
        {
            if (serializedItems == null || serializedItems.Length == 0) return null;

            return new SessionStateStoreData(SessionStateItemCollection.Deserialize(new BinaryReader(new MemoryStream(this.asCompress ? Convert.Decompress(serializedItems) : serializedItems))), SessionStateUtility.GetSessionStaticObjects(context), timeout);
        }

        //
        // SessionStateProviderBase.ReleaseItemExclusive
        //

        public override void ReleaseItemExclusive(HttpContext context, string id, object lockId)
        {
            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);

            try
            {
                Record record = client.Get(new Policy(), key, new string[] { "AN", "LockId" });
                if (record != null &&
                    Convert.ToString(record.bins["AN"]) == ApplicationName &&
                    Convert.ToInt32(record.bins["LockId"]) == Convert.ToInt32(lockId))
                {
                    // Unlock and Update Expires
                    List<Bin> bins = new List<Bin>() { new Bin("Locked", false), new Bin("Expires", DateTime.UtcNow.AddMinutes(pConfig.Timeout.TotalMinutes)) };

                    client.Put(writePolicy, key, bins.ToArray());
                }
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "ReleaseItemExclusive");
                    throw new ProviderException(exceptionMessage);
                }
                else throw e;
            }
            finally
            {
                client.Close();
            }
        }


        //
        // SessionStateProviderBase.RemoveItem
        //

        public override void RemoveItem(HttpContext context, string id, object lockId, SessionStateStoreData item)
        {
            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);

            try
            {
                Record record = client.Get(new Policy(), key, new string[] { "AN", "LockId" });
                if (record != null &&
                    Convert.ToString(record.bins["AN"]) == ApplicationName &&
                    Convert.ToInt32(record.bins["LockId"]) == Convert.ToInt32(lockId))
                {
                    client.Delete(writePolicy, key);
                }
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "RemoveItem");
                    throw new ProviderException(exceptionMessage);
                }
                else throw e;
            }
            finally
            {
                client.Close();
            }
        }



        //
        // SessionStateProviderBase.CreateUninitializedItem
        //

        public override void CreateUninitializedItem(HttpContext context, string id, int timeout)
        {
            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);

            List<Bin> bins = new List<Bin>()
            {
                new Bin("SessionId", id),
                new Bin("AN", ApplicationName),
                new Bin("Created", DateTime.UtcNow),
                new Bin("Expires", DateTime.UtcNow.AddMinutes((Double)timeout)),
                new Bin("LockDate", DateTime.UtcNow),
                new Bin("LockId", (Int32)0),
                new Bin("Timeout", timeout),
                new Bin("Locked", false),
                new Bin("SI", Value.GetAsBlob(0x0)),
                new Bin("Flags", 1),
            };

            try
            {
                client.Put(writePolicy, key, bins.ToArray());
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "CreateUninitializedItem");
                    throw new ProviderException(exceptionMessage);
                }
                else throw e;
            }
            finally
            {
                client.Close();
            }
        }


        //
        // SessionStateProviderBase.CreateNewStoreData
        //
        public override SessionStateStoreData CreateNewStoreData(HttpContext context, int timeout)
        {
            return new SessionStateStoreData(new SessionStateItemCollection(), SessionStateUtility.GetSessionStaticObjects(context), timeout);
        }



        //
        // SessionStateProviderBase.ResetItemTimeout
        //

        public override void ResetItemTimeout(HttpContext context, string id)
        {
            AerospikeClient client = new AerospikeClient(clientPolicy, asHosts);

            Key key = new Key(asNamespace, sessionSet, id);

            List<Bin> bins = new List<Bin>()
            {
                new Bin("SessionId", id),
                new Bin("AN", ApplicationName),
                new Bin("Expires", DateTime.UtcNow.AddMinutes(pConfig.Timeout.TotalMinutes)),
            };

            try
            {
                client.Put(writePolicy, key, bins.ToArray());
            }
            catch (AerospikeException e)
            {
                if (WriteExceptionsToEventLog)
                {
                    WriteToEventLog(e, "ResetItemTimeout");
                    throw new ProviderException(exceptionMessage);
                }
                else throw e;
            }
            finally
            {
                client.Close();
            }
        }


        //
        // SessionStateProviderBase.InitializeRequest
        //

        public override void InitializeRequest(HttpContext context)
        {
        }


        //
        // SessionStateProviderBase.EndRequest
        //

        public override void EndRequest(HttpContext context)
        {
        }


        //
        // WriteToEventLog
        // This is a helper function that writes exception detail to the 
        // event log. Exceptions are written to the event log as a security
        // measure to ensure private database details are not returned to 
        // browser. If a method does not return a status or Boolean
        // indicating the action succeeded or failed, the caller also 
        // throws a generic exception.
        //

        private void WriteToEventLog(Exception e, string action)
        {
            EventLog log = new EventLog();
            log.Source = eventSource;
            log.Log = eventLog;

            string message = "An exception occurred communicating with the data source.\n\n";
            message += "Action: " + action + "\n\n";
            message += "Exception: " + e.ToString();

            log.WriteEntry(message);
        }

        private List<byte[]> SplitData(byte[] data, string Key)
        {
            List<byte[]> list = new List<byte[]>();

            //list.Add(new Bin(Key, data.Length));

            byte[] buffer = null;

            for (int i = 0; i * MAX_DATA_SIZE < data.Length; i++)
            {
                buffer = new byte[Math.Min(MAX_DATA_SIZE, data.Length - (i * MAX_DATA_SIZE))];

                Array.Copy(data, i * MAX_DATA_SIZE, buffer, 0, buffer.Length);

                list.Add(buffer);
            }

            return list;
        }

        private byte[] CombineData(List<byte[]> blockData, int Size)
        {
            //int Size = Convert.ToInt32(bins[Key]);

            byte[] data = new byte[Size];

            for (int i = 0; i * MAX_DATA_SIZE < Size; i++)
            {
                blockData[i].CopyTo(data, i * MAX_DATA_SIZE);
            }

            return data;
        }
    }
}